Skip to content

[Spark 4] Streaming runner proof of concept, unbounded sources, windowed GBK and stateful ParDo on transformWithState - #39576

Draft
tkaymak wants to merge 10 commits into
apache:masterfrom
tkaymak:spark4-streaming-poc
Draft

[Spark 4] Streaming runner proof of concept, unbounded sources, windowed GBK and stateful ParDo on transformWithState#39576
tkaymak wants to merge 10 commits into
apache:masterfrom
tkaymak:spark4-streaming-poc

Conversation

@tkaymak

@tkaymak tkaymak commented Aug 1, 2026

Copy link
Copy Markdown
Contributor

Draft, opened for design discussion rather than for merging as one lump. See "How I would like to split this" at the bottom.

Addresses #36841.

What this is

The Spark 4 runner merged in #38255 is batch only. SparkStructuredStreamingRunner rejects streaming outright. This branch is a proof of concept that Spark 4's transformWithState can host the Beam streaming model, and it turns out it can, including several chained stateful operators inside a single Structured Streaming query with a watermark that propagates correctly across them.

I would like to agree on the approach before splitting this into reviewable pieces.

Approach

Dispatch seam. The shared base gains translation/PipelineTranslatorFactory.create(boolean streaming), which throws for streaming and names the Spark 4 module. runners/spark/4/src shadows that one file through the existing source override copy and returns a streaming translator instead. Every piece of transformWithState, StreamingQuery and DataSourceV2 streaming code lives only under runners/spark/4/src, because the shared base still compiles against Spark 3.5 where transformWithState does not exist. All shared base edits are behaviour preserving for Spark 3.

Source. A DataSourceV2 TableProvider wrapping any UnboundedSource. Offsets are opaque epoch counters, the driver never reads them. Executor side PartitionReaders hold UnboundedReaders in a static cache and resume from a locally cached CheckpointMark, in the style of the legacy MicrobatchSource. The row schema is payload BINARY plus eventTimestamp TIMESTAMP, so no new Catalyst encoder work was needed.

Watermark. Declared exactly once, in the read translator, on the raw rows before the typed decode. Spark forbids re-declaring it. The EventTimeWatermark plan node survives the projection and the typed map, which was the main thing I was unsure about going in. There is a test that asserts this at both the plan level and at runtime.

Stateful execution. One generic transformWithState operator, StatefulProcessor<byte[], byte[], byte[]> with Encoders.BINARY() throughout, hosting any DoFn through DoFnRunners. Stateful ParDo goes through simpleRunner plus defaultStatefulDoFnRunner, GBK through GroupAlsoByWindowViaWindowSetNewDoFn plus SystemReduceFn.buffering and KeyedWorkItems, which follows the Flink WindowDoFnOperator recipe. State is a port of the legacy SparkStateInternals onto a single MapState<String, byte[]> keyed by namespace plus tag. That layout is the only one I found that can host ReduceFnRunner's dynamically created system tags. Per @StateId column families would be faster and are an obvious follow up.

Lifecycle. StreamingEvaluationContext starts one query per leaf against the noop sink and blocks until all are terminal. cancel() reaches it through an AtomicReference set by the async translate task. Tests terminate through an idle stop listener that gracefully stops a query after N consecutive empty micro batches.

Results

Suite Tests Failures Skipped
:runners:spark:3:test 200 0 7
:runners:spark:4:test 226 0 5

121 of the Spark 4 tests are structured streaming. The streaming suite was run four separate times with identical results.

Asserted end to end: stateless ParDo, fixed and sliding window GBK, late data dropping, stateful dedup with an event time timer, chained stateful then windowed sum, and pipeline lifecycle through RUNNING, DONE and CANCELLED.

The chained case is the one that matters, and it is asserted against Spark's own StreamingQueryProgress rather than only against the computed values: one query id across every batch, exactly two transformWithStateExec operators per batch, and four strictly increasing watermarks. That rules out two separate queries that happen to sum correctly, and it rules out everything landing in one batch where all data is trivially on time.

One finding worth flagging on its own

Spark expires a transformWithState wake up at expiry <= batchWatermark. Beam fires an event time timer only once the watermark is strictly past it, and AfterWatermark.pastEndOfWindow is strict as well. Handing Beam a timer one millisecond early is therefore not merely early, it is destructive: the trigger declines, the runner is entitled to assume the timer will not be redelivered, and the on time pane is silently lost with no error anywhere. The end of window timer of a fixed window sits at exactly window.maxTimestamp(), so this hits most windows.

Fixed by bounding the fire set at min(firedExpiry, watermark - 1) and re-arming the withheld timer at firedExpiry + 1. Anyone else bridging Beam timers onto transformWithState will hit this.

Rejected at translation time, with a named message

Merging and session windows, custom triggers including AfterWatermark with early or late firings, accumulating panes, processing time timers, @OnWindowExpiration, @RequiresTimeSortedInput, side inputs on a stateful ParDo, non deterministic key coders.

Known limitations

  1. No PAssert on the streaming path. No source emits a final positive infinity watermark, so panes never finalise. Streaming tests assert against a static collector after waitUntilFinish. I would like input on what reviewers expect here, since this is the biggest departure from how the rest of Beam tests runners.
  2. Weak checkpoint recovery. Readers resume from a per JVM static cache and the source id is a fresh UUID per translation, so restart from a previous run's checkpoint is not supported. Accepted POC scope, but it is the largest gap between this and something shippable.
  3. One micro batch of timer latency. The watermark inside a stateful operator is the batch start watermark, so an end of window timer fires one micro batch after the data crossed the window end. A latency floor, not a correctness issue.
  4. stop() does not drain. An in flight micro batch finishes, anything not yet pulled is left unprocessed.
  5. Out of scope by design: session windows, custom triggers, accumulating panes, streaming side inputs, Kafka, continuous mode, portability, streaming Combine.PerKey.

Two smaller things in here

SparkSessionFactory now registers StateSchemaMetadata and MemoryWriterCommitMessage with Kryo, by name since the shared base also compiles against Spark 3. Spark 4 broadcasts a StateSchemaMetadata for every transformWithState query, so without this any stateful streaming pipeline dies on its first micro batch under spark.kryo.registrationRequired=true.

runners/spark/4/build.gradle narrows exclude "**/translation/streaming/**" to exclude "**/runners/spark/translation/streaming/**". The old glob also matched structuredstreaming/translation/streaming, so it would have silently dropped every new test in this branch from the Spark 4 module.

Not done yet, deliberately

  • Four tests in the shared base still carry @Ignore("TODO: Reactivate with streaming."). They cannot be un-ignored in place, since the shared base test tree is compiled by both modules and Spark 3 still correctly throws. Reactivating them means moving them into the Spark 4 override tree. Happy to do that, it just needs a decision on where they should live.

How I would like to split this

Assuming the approach is acceptable, five PRs, each independently green:

  1. Dispatch seam, options, lifecycle hooks, build glob fix. Shared base only, no behaviour change for Spark 3.
  2. Kryo registrations for Spark's streaming internals. Independent and self justifying.
  3. DataSourceV2 micro batch source plus tests.
  4. State and timer bridge plus unit tests. This one carries the timer fix described above and deserves its own attention.
  5. Translators, streaming evaluation context, end to end tests.

Slices 4 and 5 would benefit from a reviewer who knows Spark's Structured Streaming internals, not only Beam.

tkaymak added 10 commits July 31, 2026 11:14
Prepares the structured streaming runner for a Spark 4 streaming
translator without changing batch behavior.

- Add PipelineTranslatorFactory, the seam the Spark 4 module shadows to
  dispatch streaming pipelines. The shared base rejects streaming with a
  clear message instead of the previous generic checkArgument.
- Open up EvaluationContext (non final, protected ctor, leaves()) and add
  a no-op stop() so a streaming context can override evaluation.
- Add a createEvaluationContext hook to PipelineTranslator and skip the
  persist and lineage breaking optimizations for streaming datasets.
- Plumb the EvaluationContext into SparkStructuredStreamingPipelineResult
  so cancel() can stop a running streaming query.
- Default the state store provider to RocksDB, required by Spark 4
  transformWithState and inert for batch.
- Add watermarkDelayMillis, maxRecordsPerMicroBatch, maxBatchDurationMillis
  and streamingStopAfterIdleBatches options.
- Narrow the Spark 4 test source override exclude to the legacy DStream
  package. The previous glob also matched structuredstreaming and would
  have silently dropped the structured streaming tests.
Wraps any Beam UnboundedSource as a Spark 4 DataSourceV2 micro-batch
streaming source. The rows have exactly two columns, payload of type
BINARY holding the element encoded with a FullWindowedValueCoder, and
eventTimestamp of type TIMESTAMP holding the Beam event timestamp, so no
Catalyst encoder has to be generated for Beam types.

Offsets are opaque epoch counters. The driver never reads from the source
and never inspects its progress, latestOffset simply returns the previous
epoch plus one so Spark keeps planning micro-batches. Termination is
therefore owned by the lifecycle, not by the offsets.

The source is split once on the driver and the sub sources are memoized,
so every micro-batch plans the same stable set of partitions. On the
executor a static cache keyed by checkpoint location, source id and split
id keeps the Beam UnboundedReader alive across micro-batches, mirroring
the legacy MicrobatchSource. Checkpoint marks are cached in executor
memory only and are never committed durably, which is documented POC
scope with weak failure recovery.

UnboundedSourceDataset is the translator facing entry point. It applies
withWatermark exactly once, since Spark 4 forbids a second declaration
further down the plan.

The test confirms that the EventTimeWatermark logical node survives one
and two typed maps in both the logical and the analyzed plan, and that
the running query really tracks the watermark afterwards. It also reads a
finite set of elements end to end and checks the decoded payloads and
timestamps.
Adds a single generic transformWithState operator that can host any keyed
Beam transform, both the stateful ParDo stack and the group also by window
stack that implements a windowed GroupByKey. The mode is selected by
BeamStatefulProcessorConfig, everything else is shared.

Keys, inputs and outputs are raw Beam coder bytes and every Spark encoder
involved is BINARY or STRING, so Catalyst never has to derive a schema for a
Beam type. TwsTransformFactory owns the row layout and documents it.

Beam state is bridged by TwsStateInternals on top of BytesKV, a tiny string
to bytes store backed by exactly one Spark MapState. One map is a
requirement rather than a preference, ReduceFnRunner invents state tags at
runtime so the set of state addresses is not known when init has to declare
its state variables. Timers are bridged by TwsTimerInternals, which keeps
the full TimerData in a second map and registers only bare wake ups with
Spark, reconciling them against listTimers so an expiry can never be
registered twice.

Two boundary conditions needed care. Spark deletes the expiry it is firing
after the callback returns, so a wake up re armed at that same millisecond
is nudged one millisecond forward. Spark also expires a wake up as soon as
the expiry is at or before the batch watermark, while Beam only fires an
event time timer once the watermark is strictly past it, so timers sitting
exactly on the watermark are withheld and re armed instead of being handed
to Beam, which would decline to fire them and silently drop the on time
pane.

Processing time timers are rejected with a clear message, transformWithState
runs in a single TimeMode and this operator uses event time.

The state and timer bridges are unit tested against an in memory store, and
BeamStatefulProcessorTest runs both modes through a real streaming query on
a live SparkSession.
Shadows PipelineTranslatorFactory so streaming pipelines dispatch to a new
PipelineTranslatorStreaming, which extends PipelineTranslatorBatch to reuse
Impulse, Window.Assign, Flatten, Reshuffle, the bounded read and stateless
ParDo unchanged, while registering placeholders for the unbounded read,
GroupByKey and stateful ParDo that WS-D2 will replace. Combine.PerKey is
deliberately left unregistered so it auto-expands into GroupByKey plus ParDo.

Adds StreamingEvaluationContext, which starts one noop-sink streaming query
per leaf dataset with a checkpoint directory derived from the pipeline
options, registers an idle-stop listener for test termination, and makes
stop() idempotent and safe to call concurrently with evaluate() so cancel()
can interrupt a running streaming pipeline from another thread.
…etons

Adds StreamingTestUtils, the ListBackedUnboundedSource plus static collector
and pipeline option factories the streaming translators will be tested
against, and five skeleton test classes for the stateless ParDo, windowed
GroupByKey, stateful ParDo, chained stateful, and lifecycle scenarios. Every
skeleton compiles and runs, reporting as skipped: each is Ignore'd with a
comment naming the exact WS-D2 translator it needs once that lands.
Replaces the WS-D2 placeholders in PipelineTranslatorStreaming with real
translators for the unbounded read, GroupByKey and stateful ParDo.

ReadUnboundedTranslator wraps the Beam UnboundedSource with
UnboundedSourceDataset, which already declares the single withWatermark for
the whole query, and decodes the binary payload column into the
Dataset<WindowedValue<T>> shape every other translator consumes. The event
timestamp column is projected away, the EventTimeWatermark plan node survives
below the projection and transformWithState reads the query wide watermark
rather than a column.

GroupByKeyStreamingTranslator and StatefulParDoStreamingTranslator both encode
their input into the byte[] row layout of TwsTransformFactory, hand it to the
generic transformWithState operator in GROUP_ALSO_BY_WINDOW and STATEFUL_PARDO
mode respectively, and decode the tagged output rows back out. The stateful
ParDo splits the output by tag index and skips additional outputs that nothing
consumes, so an unused tag does not start a second streaming query.

StreamingTranslationHelpers holds the shared guards and row conversions.
Pipelines outside the POC scope are rejected at translation time with a message
naming the feature, specifically merging windows, custom triggers,
accumulating panes, processing time timers, non deterministic key coders,
non KV input to a stateful ParDo, side inputs on a stateful ParDo,
@OnWindowExpiration and @RequiresTimeSortedInput.
Turns the five Ignore'd streaming test skeletons into real, asserting
tests against the WS-D2 translators.

Every skeleton declared its pipeline as a local TestPipeline variable,
which TestPipeline itself rejects at run() because it is not a Rule
field. These tests cannot use PAssert anyway, a streaming pipeline here
never gets a final watermark so panes never finalize, and each test
needs its own options, so TestPipeline buys nothing. They now build a
plain Pipeline.create(options) and assert after waitUntilFinish against
the static collector in StreamingTestUtils. The unused TestPipeline
factory is dropped from StreamingTestUtils and the reasoning is recorded
in its javadoc.

StatelessParDoStreamingTest and StreamingPipelineLifecycleTest gain the
same relaxed Kryo SparkSessionRule as the rest of the suite, so that the
whole package shares one session and a cancelled or idle stopped query
cannot take that session down with it.

Asserted results, all derived from each test's own input list.

  Stateless ParDo, exactly 0, 2, 4 up to 18.
  Fixed 10s window count per key, a=2 and b=1.
  Sliding 10s every 5s, two a=2 panes, one per shared window.
  Late data dropped, a=1 and sentinel=1, never a=2.
  Stateful ParDo, a, b, c once each plus two timer sentinels.
  Chained dedup into windowed sum, a=8 and b=10.

The chained case is the one that matters, a=13 would mean the dedup
state was lost and an empty result would mean the watermark got stuck
between the two operators, so exactly a=8 and b=10 is the observable
proof of cross operator watermark propagation.

The late data test is the only one that depends on how the source round
robins elements across splits, so it pins maxRecordsPerMicroBatch to 1,
spells out the resulting per batch schedule in a comment, and asserts
the split count assumption up front rather than leaving it silent.

The lifecycle tests wait for a query to actually be active before
cancelling, since translation is asynchronous and an early cancel would
find a null evaluation context, and they check the query really stopped
rather than only that the state was relabelled.
Spark 4 broadcasts a StateSchemaMetadata to the executors for every
transformWithState query, using the user Kryo instance. Nothing
registered that class, so any Beam streaming pipeline with state or
timers died on its first micro-batch as soon as
spark.kryo.registrationRequired was on, which is the default for this
module's tests and a perfectly reasonable production setting. The same
applies to MemoryWriterCommitMessage, the commit message of Spark's
memory sink, which is nested inside the already registered
DataWritingSparkTaskResult.

Both are registered by name, because the shared runner base also
compiles against Spark 3 where neither class exists, and both are
registered with a JavaSerializer rather than Kryo's default field
serializer. Both are Scala case classes holding further Scala and Spark
types, immutable.Map, StructType, avro Schema and Row, none of which are
registered either, so going through Java serialization covers the whole
object graph at once instead of forcing this list to track Spark's
internal field layout. Neither object is on a hot path.

Every streaming test now runs with the module default of
spark.kryo.registrationRequired=true, the per test relaxations are gone.
SparkKryoRegistratorStreamingTest names both classes at compile time
against the Spark 4 classpath, so a future rename becomes a compile
error rather than a silently dropped registration.
Two javadoc paragraphs in the shared base drifted from google-java-format,
one from the streaming note added to SparkStructuredStreamingRunner and one
from the new Kryo registrations. The shared base sources under
runners/spark/src are checked by the :runners:spark project rather than by
:runners:spark:3 or :runners:spark:4, whose own spotless targets only see
their per version override directories, so it is easy to miss.
ChainedStatefulStreamingTest asserts what the chained pipeline computes.
It cannot distinguish a single query holding two transformWithState
operators from two queries that happen to add up to the same numbers, and
it cannot show whether the watermark moved at all or whether the whole
input simply landed in one micro-batch where everything is trivially on
time.

ChainedStatefulStreamingEvidenceTest asserts the run itself, against
Spark's own StreamingQueryProgress. It pins one record per split per
micro-batch so the watermark has to climb in steps, records every
progress event through a listener, and then asserts that all of them
carry one query id, that one micro-batch reports exactly two
transformWithStateExec state operators, and that the watermark takes at
least three strictly increasing values.

The second test adds a tap between the two operators, so the late record
is observed leaving the dedup operator and absent from the windowed sum.
That is the difference between a record excluded for lateness and a
record that never arrived.

Both tests print the raw per batch progress they recorded, since that
output is the evidence the phase gate report quotes.

Also drops four javadoc references to a Kryo relaxation that no longer
exists anywhere in the suite.
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant